#include "Synchro.hpp"
#include <cstring>

#define COMPUTE_NEW_ARG \
	if (new_mi) { \
		computeNewArg(mi); \
	}

Synchro::Synchro(AlgorithmRef algEngine) :
	algEngine(algEngine) {

	allocStructures();
}

/** default destructor */
Synchro::~Synchro() {
	deallocStructures();
}

/** Method to return some info about the nlp */
bool Synchro::get_nlp_info(Index& n, Index& m, Index& nnz_jac_g,
		Index& nnz_h_lag, IndexStyleEnum& index_style) {
	n = algEngine.subProblemsCount + 1;
	m = algEngine.restrictions + 1;
	nnz_jac_g = n * m; //very dense
	nnz_h_lag = 0;
	index_style = C_STYLE;

	return true;
}

/** Method to return the bounds for my problem */
bool Synchro::get_bounds_info(Index n, Number* mi_l, Number* mi_u, Index m,
		Number* g_l, Number* g_u) {
	for (int i = 0; i < n; ++i) {
		mi_l[i] = -1e19;
		mi_u[i] = 1e19;
	}

	if (m > 1) {
		//problem constraints
		for (int i = 0; i < m - 1; ++i) {
			g_l[i] = algEngine.g_l[i];
			g_u[i] = algEngine.g_u[i];
		}
	}

	//normalizacja mi
	g_l[m - 1] = 1;
	g_u[m - 1] = 1;

	return true;
}

/** Method to return the starting point for the algorithm */
bool Synchro::get_starting_point(Index n, bool init_x, Number* x, bool init_z,
		Number* z_L, Number* z_U, Index m, bool init_lambda, Number* lambda) {
	assert(init_z == false);
	assert(init_lambda == false);

	if (init_x) {
		for (Index i = 0; i < n; ++i) {
			x[i] = 1.0 / n;
		}
	}

	return true;
}

/** Method to return the objective value */
bool Synchro::eval_f(Index n, const Number* mi, bool new_mi, Number& obj_value) {
	COMPUTE_NEW_ARG;

	algEngine.problemDefinition.eval_f(algEngine.dim, xOrgSpace, new_mi,
			obj_value);

	return true;
}

/** Method to return the gradient of the objective */
bool Synchro::eval_grad_f(Index n, const Number* mi, bool new_mi,
		Number* grad_mi) {
	COMPUTE_NEW_ARG;

	//funGradient - gradient oryginalnej funkcji
	algEngine.problemDefinition.eval_grad_f(algEngine.dim, xOrgSpace, new_mi,
			orgGradient);

	Number sum = 0;
	for (int i = 0; i < algEngine.dim; ++i) {
		sum += orgGradient[i] * algEngine.lastOptX[i];
	}
	grad_mi[0] = sum;

	for (int p = 0; p < algEngine.subProblemsCount; ++p) {
		sum = 0;
		for (int i = 0; i < algEngine.dim; ++i) {
			sum += orgGradient[i] * algEngine.subProblems[p]->xOrgSpace[i];
		}
		grad_mi[p + 1] = sum;
	}

	return true;
}

/** Method to return the constraint residuals */
bool Synchro::eval_g(Index n, const Number* mi, bool new_mi, Index m, Number* g) {
	COMPUTE_NEW_ARG;

	bool ret = true;
	if (m > 1) {
		//problem constraints
		ret = algEngine.problemDefinition.eval_g(algEngine.dim, xOrgSpace,
				true, m - 1, g);
	}

	//mi normalization
	Number sum = 0;
	for (int i = 0; i < n; ++i) {
		sum += mi[i];
	}
	g[m - 1] = sum;

	return ret;
}

/** Method to return:
 *   1) The structure of the jacobian (if "values" is NULL)
 *   2) The values of the jacobian (if "values" is not NULL)
 */
bool Synchro::eval_jac_g(Index n, const Number* mi, bool new_mi, Index m,
		Index nele_jac, Index* iRow, Index *jCol, Number* values) {
	COMPUTE_NEW_ARG;

	bool ret = true;
	if (values != NULL) {
		if (m > 1) {
			//problem constraints
			ret = algEngine.problemDefinition.eval_jac_g(algEngine.dim,
					xOrgSpace, new_mi, m - 1, algEngine.orgNnzJac,
					algEngine.jacIRow, algEngine.jacJCol, orgJacob);
			memset(values, 0, sizeof(Number) * nele_jac);

			for (Index orgNnzInd = 0; orgNnzInd < algEngine.orgNnzJac; orgNnzInd++) {
				Index restrNum = algEngine.jacIRow[orgNnzInd];
				Index varNum = algEngine.jacJCol[orgNnzInd];
				values[n * restrNum] += orgJacob[orgNnzInd]
						* algEngine.lastOptX[varNum];
				for (Index pp1 = 1; pp1 < n; ++pp1) {
					values[n * restrNum + pp1] += orgJacob[orgNnzInd]
							* algEngine.subProblems[pp1 - 1]->xOrgSpace[varNum];
				}
			}
		}

		//mi normalization
		for (Index i = 0; i < n; ++i) {
			values[n * (m - 1) + i] = 1;
		}
	} else {
		//struktura
		int ind = 0;
		for (int i = 0; i < m; ++i) {
			for (int j = 0; j < n; ++j) {
				iRow[ind] = i;
				jCol[ind] = j;
				ind++;
			}
		}
		assert(ind == nele_jac);
	}
	return ret;
}

/** Method to return:
 *   1) The structure of the hessian of the lagrangian (if "values" is NULL)
 *   2) The values of the hessian of the lagrangian (if "values" is not NULL)
 */
bool Synchro::eval_h(Index n, const Number* mi, bool new_mi, Number obj_factor,
		Index m, const Number* lambda, bool new_lambda, Index nele_hess,
		Index* iRow, Index* jCol, Number* values) {
	COMPUTE_NEW_ARG;

	return false;
}

/** This method is called when the algorithm is complete so the TNLP can store/write the solution */
void Synchro::finalize_solution(SolverReturn status, Index n, const Number* mi,
		const Number* z_L, const Number* z_U, Index m, const Number* g,
		const Number* lambda, Number obj_value IPOPT8_FINARGS) {
	computeNewArg(mi);
}

void Synchro::computeNewArg(const Number* mi) {
	for (int i = 0; i < algEngine.dim; ++i) {
		Number vecEl = mi[0] * algEngine.lastOptX[i];
		for (int p = 0; p < algEngine.subProblemsCount; ++p) {
			vecEl += mi[p + 1] * algEngine.subProblems[p]->xOrgSpace[i];
		}
		xOrgSpace[i] = vecEl;
	}
}

void Synchro::allocStructures() {
	//original function cache
	Number* pData = new Number[2 * algEngine.dim + algEngine.orgNnzJac];

	orgGradient = pData;
	xOrgSpace = orgGradient + algEngine.dim;
	orgJacob = xOrgSpace + algEngine.dim;
}

void Synchro::deallocStructures() {
	if (orgGradient != NULL) {
		delete[] orgGradient;
		orgGradient = NULL;
		xOrgSpace = NULL;
		orgJacob = NULL;
	}
}
